import pefile
import shutil
import hashlib
import os
import numpy as np
import array
import cv2
from basicTools.myTLSH import MyTlsh


class FileProcess:
    """
    处理文件夹中的内容
    1. 判断文件类型
    2. 文件转md5命名
    3. 文件生成灰度图
    4. 文件生成TLSH值
    """

    def __init__(self):
        self.myTlsh = MyTlsh()

    def moveAndRenameByMd5(self, originPath, savePath, sampleName):
        """
        将样本改成用md5命名，并返回其md5值
        input:
            origginPath:原始样本的路径
            savePath:期望生成样本的父路径
            sampleName:修改的样本名
        """
        filePath = os.path.join(originPath, sampleName)
        tarFilePath = os.path.join(savePath, sampleName)
        # 复制文件
        shutil.copy2(filePath, tarFilePath)
        # 重命名
        fileMd5 = self.calculateSampleMd5(tarFilePath)
        os.rename(tarFilePath, os.path.join(savePath, fileMd5))
        return fileMd5

    def calculateMd5(self, savePath, sampleName):
        tarFilePath = os.path.join(savePath, sampleName)
        fileMd5 = self.calculateSampleMd5(tarFilePath)
        return fileMd5

    def calculateSampleMd5(self, filePath, block_size=65536):
        """
        计算文件的MD5值
        input: filePath: 样本路径
        """
        hasher = hashlib.md5()
        with open(filePath, "rb") as f:
            for block in iter(lambda: f.read(block_size), b""):
                hasher.update(block)
        return hasher.hexdigest()

    def isPeFile(self, file_path):
        """
        判断文件是否为pe样本
        """
        try:
            pe = pefile.PE(file_path)
            return True
        except pefile.PEFormatError:
            return False

    def generateGreyImage(self, input_dir, output_dir, filename, str=""):
        """
        生成二进制灰度图片
        """
        if str != "":
            temp = str.split("_")[1]
            output_dir += temp
            output_dir += "/"
        out_file = os.path.splitext(str + os.path.basename(filename))[0] + ".png"
        out_file_full = output_dir + out_file
        input_file_path = os.path.join(input_dir, filename)

        f = open(input_file_path, "rb")
        ln = os.path.getsize(input_file_path)
        width = 256
        rem = ln % width

        a = array.array("B")
        a.fromfile(f, ln - rem)
        f.close()
        g = np.reshape(a, (len(a) // width, width))
        g = np.uint8(g)
        img_resized = cv2.resize(g, (105, 105), interpolation=cv2.INTER_CUBIC)
        cv2.imwrite(out_file_full, img_resized)

    def generateTLSH(self, filePath):
        """
        input: sample path
        output: tlsh value
        """
        return self.myTlsh.calculateTlsh(filePath)


if __name__ == "__main__":
    fileProcess = FileProcess()
    fileProcess.generateGreyImage(
        r"D:/HGMSim/HGMSimBackEndByFlask/data/md5Samples",
        r"D:/HGMSim/HGMSimBackEndByFlask/data/greyImage/",
        "0db3ab52940c630a9f73691862b45758",
    )
